import os
from src.common import commonUtils
from pydantic import BaseModel, Field
from langgraph.graph import StateGraph
from src.common.logger import getLogger
from typing import TypedDict, List, Dict
from langgraph.constants import START, END
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

logger = getLogger()

class Plan(BaseModel):
    index: int = Field(description = "计划序号")
    detail: str = Field(description = "计划详情描述")
    funcName: str = Field(description = "可用的工具名")
    funcParam: Dict[str, str] = Field(description = "json格式的工具的输入参数")

class PlanList(BaseModel):
    plans: List[Plan] = Field(description = "Plan对象的列表")

class SelfDiscoverState(TypedDict):
    query: str
    modules: str
    instructions: str
    plans: List[Plan]
    answer: str

class SelfDiscoverAgent:

    def __init__(self, llm_model, agent_tools):
        self.llm_model = llm_model
        self.agent_tools = agent_tools

    def select_node(self, state: SelfDiscoverState):
        logger.info("SelfDiscoverAgent select_node start")
        template = """
            你是一个推理策略选择器。请根据用户问题，从以下预定义推理模块中选择最相关的若干项（至少1项，最多3项）：

            可用模块：
            - Step-by-step reasoning（逐步推理）
            - Analogy（类比推理）
            - Causal reasoning（因果推理）
            - Counterfactual thinking（反事实推理）
            - Abstraction and generalization（抽象与泛化）
            - Decomposition（任务分解）
            - Example-driven reasoning（示例驱动）
            - Constraint satisfaction（约束满足）
            
            要求：
            - 仅输出选中的模块名称，每行一个，不要解释。
            - 必须严格使用上述模块名称，不得改写或新增。
            
            用户问题：{question}
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        response = chain.invoke({ "question": state["query"] })
        logger.info(f"SelfDiscoverAgent select_node response: {response}")
        return { "modules": response }

    def adapt_node(self, state: SelfDiscoverState):
        logger.info("SelfDiscoverAgent adapt_node start")
        template = """
            你是一个推理策略适配器。请根据用户问题和选定的推理模块，为每个模块生成一段定制化的推理指令（instruction）。每条指令应具体、可操作，并贴合问题上下文。
            
            要求：
            - 每个模块对应一条指令；
            - 指令以“请...”开头，聚焦如何应用该模块解决当前问题；
            - 不要重复问题，不要输出答案，只输出适配后的指令；
            - 每条指令单独一行。
            
            用户问题：{question}
            选定模块：{modules}
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        response = chain.invoke({ "question": state["query"], "modules": state["modules"] })
        logger.info(f"SelfDiscoverAgent adapt_node response: {response}")
        return { "instructions": response }

    def plan_node(self, state: SelfDiscoverState):
        logger.info("SelfDiscoverAgent plan_node start")
        template = """
            你是一个推理结构编排器。请参考提供的可用工具将以下适配后的推理指令整合成一个清晰、有序的多步推理计划。每一步应编号，并保持逻辑连贯。
            
            要求：
            - 输出格式：Step 1: ...\nStep 2: ...\n...
            - 不要添加额外说明或总结；
            - 必须覆盖所有输入指令，可合并相似步骤；
            - 步骤应可被后续执行。
            
            适配后的指令：{instructions}
            可用的工具：{tool_descs}
        """
        tool_descs = commonUtils.build_tools_description(self.agent_tools)
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model.with_structured_output(PlanList)
        response = chain.invoke({ "question": state["query"], "instructions": state["instructions"], "tool_descs": tool_descs })
        logger.info(f"SelfDiscoverAgent plan_node response: \n{response}")
        return { "plans": response.plans }

    def solve_node(self, state: SelfDiscoverState):
        logger.info("SelfDiscoverAgent solve_node start")

        results = []
        detail_plans = []
        plans = state["plans"]
        for plan in plans:
            logger.info(f"SelfDiscoverAgent solve_node plan: \n{plan}")
            detail_plans.append(plan.detail)
            if plan.funcName:
                for tool in self.agent_tools:
                    if tool.name == plan.funcName:
                        result = tool.func(*plan.funcParam)
                        results.append(result)

        context = "\n".join(results)
        structured_plan = "\n".join(detail_plans)

        template = """
            你是一个推理执行器。请严格按照提供的上下文和推理计划，一步步解答用户问题。在每一步中展示你的思考过程，最后给出简洁、准确的最终答案。
            
            要求：
            - 先复述推理计划；
            - 然后按步骤执行，每步标明“Step X: ...”；
            - 最终答案以“Final Answer: ”开头，且仅包含答案内容；
            - 不得跳过步骤，不得引入计划外的推理。
            
            上下文：{context}
            用户问题：{question}
            推理计划：{structured_plan}
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        response = chain.invoke({ "context": context, "question": state["query"], "structured_plan": structured_plan })
        logger.info(f"SelfDiscoverAgent solve_node response len: {len(response)}")
        return { "answer": response }

    def build_graph(self):
        logger.info("SelfDiscoverAgent build_graph start")
        graph = StateGraph(SelfDiscoverState)
        graph.add_node("select", self.select_node)
        graph.add_node("adapt", self.adapt_node)
        graph.add_node("plan", self.plan_node)
        graph.add_node("solve", self.solve_node)

        graph.add_edge(START, "select")
        graph.add_edge("select", "adapt")
        graph.add_edge("adapt", "plan")
        graph.add_edge("plan", "solve")
        graph.add_edge("solve", END)

        workflow = graph.compile()

        save_path = "D:/Downloads/taixu/images/agentics"
        image_path = os.path.join(save_path, "Agent_SelfDiscover_Workflow.png")
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        with open(image_path, 'wb') as file:
            file.write(workflow.get_graph().draw_mermaid_png())

        return workflow

    def invoke(self, query):
        logger.info(f"SelfDiscoverAgent invoke query: {query}")
        workflow = self.build_graph()
        response = workflow.invoke({ "query": query })
        answer = response.get("answer", None)
        logger.info(f"SelfDiscoverAgent invoke answer len: {len(str(answer))}")
        return { "chain_result": answer }
